import os
import argparse

def extract_full_aladdin_summary(summary_path):
    """Extract all relevant fields from Aladdin summary for comparative table."""
    result = {
        "Cycle（执行周期）": "",
        "Avg Power（平均功耗）": "",
        "Idle FU Cycles": "",
        "Avg FU Power": "",
        "Avg FU Dynamic Power": "",
        "Avg FU Leakage Power": "",
        "Avg MEM Power": "",
        "Avg MEM Dynamic Power": "",
        "Avg MEM Leakage Power": "",
        "Total Area": "",
        "FU Area": "",
        "MEM Area": "",
        "Num of Multipliers (32-bit)": "",
        "Num of Adders (32-bit)": "",
        "Num of Shifters (32-bit)": "",
        "Num of Registers (32-bit)": ""
    }

    if not os.path.isfile(summary_path):
        return None

    with open(summary_path, "r") as f:
        for line in f:
            if "Upsampled" in line:
                continue
            for key in result:
                if key.startswith("Cycle") and "Cycle :" in line:
                    print(line)
                    result[key] = line.split(":")[1].strip().replace(" cycles", "")
                elif key.startswith("Avg Power") and "Avg Power:" in line:
                    result[key] = line.split(":")[1].strip()
                elif key in line:
                    result[key] = line.split(":")[1].strip()

    return result


def generate_comparative_markdown_table(parent_dir):
    """Generate a full comparative markdown table from all subdirectories."""
    entries = sorted([d for d in os.listdir(parent_dir) if os.path.isdir(os.path.join(parent_dir, d))])
    data_by_dir = {}

    for entry in entries:
        summary_path = os.path.join(parent_dir, entry, "stencil_summary")
        parsed = extract_full_aladdin_summary(summary_path)
        if parsed:
            data_by_dir[entry] = parsed

    # Determine row headers
    if not data_by_dir:
        return "No valid summaries found."

    row_keys = list(next(iter(data_by_dir.values())).keys())
    col_headers = list(data_by_dir.keys())

    # Construct markdown table
    markdown = []
    header = "| 指标 | " + " | ".join(col_headers) + " |"
    divider = "|---" * (len(col_headers) + 1) + "|"
    markdown.append(header)
    markdown.append(divider)

    for key in row_keys:
        row = f"| {key} | " + " | ".join(data_by_dir[dir_name][key] for dir_name in col_headers) + " |"
        markdown.append(row)

    return "\n".join(markdown)


def write_summary_markdown_to_file(parent_dir, output_file="summary.md"):
    markdown = generate_comparative_markdown_table(parent_dir)
    with open(output_file, "w", encoding="utf-8") as f:
        f.write(markdown)
    return output_file

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate comparative markdown table from Aladdin summaries.")
    parser.add_argument("dir", help="Directory containing subdirectories with stencil_summary files.")
    parser.add_argument("--output", default="summary.md", help="Output markdown file name (default: summary.md)")
    args = parser.parse_args()

    output_path = write_summary_markdown_to_file(args.dir, args.output)
    print(f"Markdown summary written to: {output_path}")
